{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from netCDF4 import Dataset\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "from matplotlib.cm import get_cmap\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "import matplotlib.gridspec as gridspec\n",
    "from wrf import getvar, interplevel, to_np, get_basemap, latlon_coords, destagger\n",
    "import glob\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import cmocean\n",
    "import pyresample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = \"mynn\"\n",
    "dom1 = \"d01\"\n",
    "dom2 = \"d02\"\n",
    "active_dom = \"d02\"\n",
    "sens = \"no_fractional_seaice\"\n",
    "day = \"12\"\n",
    "#root = \"/glade/scratch/tjuliano/doe_comble/new_lam_runs/WRF/WRF-ASR-CAO/test/\"\n",
    "root = \"/glade/campaign/uwyo/wyom0122/lam/new_runs/\"\n",
    "cases = [\"031220_no_edmf\",\"031220_mynn_l3\",\"031220\",\"031220_ysu\"]\n",
    "rst = \"/RESTART_D\"\n",
    "if active_dom == \"d01\":\n",
    "    plt_titles = [r\"$\\Delta$x=3km, level 2.5\", \"$\\Delta$x=3km, level 3\", \"$\\Delta$x=3km, level 2.5 EDMF\", \"$\\Delta$x=3km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=3km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=3km\"\"\\n\"\"Level 3\", \"$\\Delta$x=3km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=3km\"\"\\n\"\"YSU\"]\n",
    "    dx = 3000.\n",
    "else:\n",
    "    plt_titles = [r\"$\\Delta$x=1km, level 2.5\", \"$\\Delta$x=1km, level 3\", \"$\\Delta$x=1km, level 2.5 EDMF\", \"$\\Delta$x=1km, YSU\"]\n",
    "    plt_titles2 = [r\"$\\Delta$x=1km\"\"\\n\"\"Level 2.5\", \"$\\Delta$x=1km\"\"\\n\"\"Level 3\", \"$\\Delta$x=1km\"\"\\n\"\"Level 2.5 EDMF\", \"$\\Delta$x=1km\"\"\\n\"\"YSU\"]\n",
    "    dx = 1000.\n",
    "plt_titles = [\"Level 2.5\", \"Level 3\", \"Level 2.5 EDMF\", \"YSU\"]\n",
    "savedir = \"figs_new_nz136/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "lwprng = np.arange(-30,80,5)\n",
    "times = [9,11,13,19,25,31,37,43,49]\n",
    "time_str = ['28_1030z','28_1130z','28_1230z','28_1530z','28_1830z','28_2130z','29_0030z','29_0330z','29_0630z']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "lev = 10\n",
    "skipx = 6\n",
    "skipy = 6\n",
    "skipx2 = skipx*3\n",
    "skipy2 = skipy*3\n",
    "scale = 50\n",
    "barbw = 0.04\n",
    "g = 9.81"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "s1 = Dataset(root + cases[2] + \"/wrfinput_\" + dom1 + \"_trim\")\n",
    "cosalpha = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()\n",
    "\n",
    "s1 = Dataset(root + cases[2] + \"/wrfinput_\" + dom2 + \"_trim\")\n",
    "cosalpha2 = s1.variables['COSALPHA'][-1,:,:]\n",
    "sinalpha2 = s1.variables['SINALPHA'][-1,:,:]\n",
    "s1.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = 'cells'\n",
    "if region == 'cells':\n",
    "###### OPEN CELLS\n",
    "#\tlon_arr = [12.0, 12.0, 16.0, 16.0]\n",
    "#\tlat_arr = [69.5, 71.5, 71.5, 69.5]\n",
    "\tlon_arr = [12.5, 12.5, 16.5, 16.5]\n",
    "\tlat_arr = [69.0, 71.0, 71.0, 69.0]\n",
    "elif region == 'rolls':\n",
    "###### ROLLS/CLOSED CELLS\n",
    "##        lon_arr = [7.5,  7.5, 12.5, 12.5]\n",
    "##        lat_arr = [71.5, 73.0, 73.0, 71.5]\n",
    "\tlon_arr = [8.0, 8.0, 12.0, 12.0]\n",
    "\tlat_arr = [74.5, 76.5, 76.5, 74.5]\n",
    "else:\n",
    "        print ('You are an idiot!')\n",
    "        \n",
    "#lat2 = getvar(s1,\"lat\",meta=False)\n",
    "#lon2 = getvar(s1,\"lon\",meta=False)\n",
    "#hgt = getvar(s1,\"HGT\",meta=False)\n",
    "#lat_d01 = lat2\n",
    "#lon_d01 = lon2\n",
    "#hgt_d01 = hgt\n",
    "ll_lon = lon_arr[0]\n",
    "ll_lat = lat_arr[0]\n",
    "ul_lon = lon_arr[0]\n",
    "ul_lat = lat_arr[1]\n",
    "ur_lon = lon_arr[2]\n",
    "ur_lat = lat_arr[1]\n",
    "lr_lon = lon_arr[2]\n",
    "lr_lat = lat_arr[0]\n",
    "# center point of overall domain\n",
    "ref_lat = (ll_lat + ur_lat)/2.\n",
    "#ref_lon = -(abs(ll_lon) + abs(ur_lon))/2.\n",
    "ref_lon = (ll_lon + ur_lon)/2.\n",
    "######################################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files_wind0 = sorted(glob.glob(root + cases[0] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind1 = sorted(glob.glob(root + cases[1] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind2 = sorted(glob.glob(root + cases[2] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "files_wind3 = sorted(glob.glob(root + cases[3] + rst + \"/wrfout_wind_\" + active_dom + \"*\"))\n",
    "\n",
    "times = [12]\n",
    "first_pass1 = 0\n",
    "for i in times: #np.arange(12,len(files_cloud0),6):\n",
    "\tfor j in np.arange(len(cases)):\n",
    "\t\tprint ('Reading file...')\n",
    "\t\tif j == 0:\n",
    "\t\t\ts1 = Dataset(files_wind0[i])\n",
    "\t\t\tprint (files_wind0[i])\n",
    "\t\telif j == 1:\n",
    "\t\t\ts1 = Dataset(files_wind1[i])\n",
    "\t\t\tprint (files_wind1[i])\n",
    "\t\telif j == 2:\n",
    "\t\t\ts1 = Dataset(files_wind2[i])\n",
    "\t\t\tprint (files_wind2[i])\n",
    "\t\telif j == 3:\n",
    "\t\t\ts1 = Dataset(files_wind3[i])\n",
    "\t\t\tprint (files_wind3[i])\n",
    "\t\tif j == 0:\n",
    "\t\t\ts3 = Dataset(root + cases[0] + \"/wrfinput_\" + active_dom)\n",
    "\t\t\tlat = getvar(s3,\"lat\",meta=False)\n",
    "\t\t\tlon = getvar(s3,\"lon\",meta=False)\n",
    "\t\tseaice = getvar(s3,\"SEAICE\",meta=False)\n",
    "\t\tu = getvar(s1,\"ua\",units=\"m s-1\",meta=False)\n",
    "\t\tv = getvar(s1,\"va\",units=\"m s-1\",meta=False)\n",
    "\t\tw = getvar(s1,\"wa\",units=\"m s-1\",meta=False)\n",
    "\t\tt2 = getvar(s1,\"T2\",meta=False)\n",
    "\t\tz = getvar(s1,\"z\", units=\"m\",meta=False)\n",
    "\t\tpblh = getvar(s1,\"PBLH\",meta=False)\n",
    "\n",
    "\t\tif j == 0:\n",
    "\t\t\tfig = plt.figure(figsize=(18,15))\n",
    "\t\t\tax1 = fig.add_axes([0.9,0.675,0.02,0.3])\n",
    "\t\t\tax2 = fig.add_axes([0.9,0.35,0.02,0.3])\n",
    "\t\t\tax3 = fig.add_axes([0.9,0.025,0.02,0.3])\n",
    "\n",
    "\t\tprint ('Interpolating...')\n",
    "\t\thgt = pblh/2.\n",
    "\t\thgt2 = pblh/10.\n",
    "\t\thgt2 = 100.\n",
    "\t\tui = interplevel(u, z, hgt2, meta=False)\n",
    "\t\tvi = interplevel(v, z, hgt2, meta=False)\n",
    "\t\twi = interplevel(w, z, hgt)\n",
    "\n",
    "\t\tif active_dom == \"d01\":\n",
    "\t\t\tuearth = ui*cosalpha - vi*sinalpha\n",
    "\t\t\tvearth = vi*cosalpha + ui*sinalpha\n",
    "\t\telse:\n",
    "\t\t\tuearth = ui*cosalpha2 - vi*sinalpha2\n",
    "\t\t\tvearth = vi*cosalpha2 + ui*sinalpha2\n",
    "\n",
    "\t\tuearth2 = destagger(uearth,0)\n",
    "\t\tvearth2 = destagger(vearth,1)\n",
    "\t\tconv = (uearth2[:,1:]-uearth2[:,0:-1])/dx + (vearth2[1:,:]-vearth2[0:-1,:])/dx\n",
    "        \n",
    "\t\tif first_pass1 == 0:\n",
    "\t\t\tlon2d = s1.variables['XLONG'][-1,:,:]\n",
    "\t\t\tlat2d = s1.variables['XLAT'][-1,:,:]\n",
    "\t\t\tgrid = pyresample.geometry.GridDefinition(lats=lat2d, lons=lon2d)\n",
    "\n",
    "            # Define points\n",
    "\t\t\tswath = pyresample.geometry.SwathDefinition(lons=lon_arr, lats=lat_arr)\n",
    "            # Determine nearest (w.r.t. great circle distance) neighbour in the grid.\n",
    "\t\t\t_, _, index_array, distance_array = pyresample.kd_tree.get_neighbour_info(\n",
    "\t\t\tsource_geo_def=grid, target_geo_def=swath, radius_of_influence=5000,neighbours=1)\n",
    "            # get_neighbour_info() returns indices in the flattened lat/lon grid. Compute the 2D grid indices:\n",
    "\t\t\tcalipso_j1, calipso_i1 = np.unravel_index(index_array, grid.shape)\n",
    "\t\t\tprint (calipso_j1, calipso_i1)\n",
    "\n",
    "\t\t\tfirst_pass1 = 1\n",
    "\n",
    "\t\tcalipso_j = calipso_j1\n",
    "\t\tcalipso_i = calipso_i1\n",
    "        \n",
    "\t\tconv2 = conv[calipso_j[0]:calipso_j[2],calipso_i[1]:calipso_i[3]]\n",
    "\n",
    "#\t\tup = uearth - np.mean(uearth.ravel())#\n",
    "#\t\tvp = vearth - np.mean(vearth.ravel())\n",
    "\n",
    "\t################################################################################################################\n",
    "\t############################################### LWP ############################################################\n",
    "\t################################################################################################################\n",
    "\t\tprint ('Plotting LWP...')\n",
    "        \n",
    "\t\tplt.subplot(3,4,j+1)\n",
    "\t\tm = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='h')\n",
    "\t\tx, y = m(to_np(lon), to_np(lat))\n",
    "\t\t# Make the contour plot\n",
    "\t\tnorm = mpl.colors.Normalize(vmin=0, vmax=3000)\n",
    "\t\trng = np.arange(-0.8,0.9,0.1)\n",
    "\t\tw_contourf = m.contourf(x,y,to_np(wi[:,:]),cmap=cmocean.cm.oxy,levels=rng,extend='both',zorder=2)\n",
    "\t\tm.contour(x,y,to_np(seaice),levels=[0.0001,0.25,0.5,0.75,1.0],colors=['magenta','pink','pink','pink'],linewidths=2.,zorder=4)\n",
    "\n",
    "\t\t# Add geographic outlines\n",
    "\t\tm.drawcoastlines(linewidth=2.5,zorder=3)\n",
    "\t\tif j == 0:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[True,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\telse:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\tm.drawcountries(linewidth=2.0,zorder=3)\n",
    "\t\tm.fillcontinents(color='coral',lake_color='coral',zorder=3)\n",
    "        \n",
    "\t\tx_anx,y_anx = m(15.874211,69.232286)\n",
    "\t\tm.scatter(x_anx,y_anx,marker='D',s=360,facecolor='yellow',edgecolor='k',zorder=8)\n",
    "        \n",
    "\t\tplt.title(plt_titles[j],fontdict = {'fontsize' : 22})\n",
    "\t\tif j == 1:\n",
    "\t\t\tplt.suptitle(files_wind0[i][-19:],y=1.02,fontsize=28)\n",
    "            \n",
    "\t\tif j == 0 or j == 2:\n",
    "\t\t\tif j == 0:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.5)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.5)\n",
    "\t\t\telse:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.4)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.4)\n",
    "\t\t\tm.plot([x1,x2],[y1,y2],c='limegreen',lw=3.,ls='--',zorder=8)\n",
    "\n",
    "\n",
    "\t\tplt.subplot(3,4,4+j+1)\n",
    "\t\tm = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='h')\n",
    "\t\tx, y = m(to_np(lon), to_np(lat))\n",
    "\t\t# Make the contour plot\n",
    "\t\tnorm = mpl.colors.Normalize(vmin=0, vmax=3000)\n",
    "\t\trng = [-0.003,-0.0025,-0.002,-0.0015,-0.001,-0.0005,0,0.0005,0.001,0.0015,0.002,0.0025,0.003]\n",
    "\t\tconv_contourf = m.contourf(x[1:,1:],y[1:,1:],to_np(conv[:,:]),levels=rng,cmap=cmocean.cm.balance,extend='both',zorder=2)\n",
    "\t\tw_contours = m.contour(x,y,to_np(wi[:,:]),colors='magenta',levels=[0.25],zorder=3)\n",
    "\t\tm.contour(x,y,to_np(seaice),levels=[0.0001,0.25,0.5,0.75,1.0],colors=['magenta','pink','pink','pink'],linewidths=2.,zorder=4)\n",
    "\n",
    "\t\t# Add geographic outlines\n",
    "\t\tm.drawcoastlines(linewidth=2.5,zorder=3)\n",
    "\t\tif j == 0:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[True,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\telse:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\tm.drawcountries(linewidth=2.0,zorder=3)\n",
    "\t\tm.fillcontinents(color='coral',lake_color='coral',zorder=3)\n",
    "        \n",
    "\t\tx_anx,y_anx = m(15.874211,69.232286)\n",
    "\t\tm.scatter(x_anx,y_anx,marker='D',s=360,facecolor='yellow',edgecolor='k',zorder=8)\n",
    "        \n",
    "\t\tif j == 0 or j == 2:\n",
    "\t\t\tif j == 0:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.5)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.5)\n",
    "\t\t\telse:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.4)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.4)\n",
    "\t\t\tm.plot([x1,x2],[y1,y2],c='limegreen',lw=3.,ls='--',zorder=8)\n",
    "        \n",
    "        \n",
    "\t\tplt.subplot(3,4,8+j+1)\n",
    "\t\tm = Basemap(llcrnrlon=ll_lon,llcrnrlat=ll_lat,urcrnrlon=ur_lon,urcrnrlat=ur_lat,lon_0=ref_lon,lat_0=ref_lat,projection='cass',resolution='h')\n",
    "\t\tx, y = m(to_np(lon), to_np(lat))\n",
    "\t\t# Make the contour plot\n",
    "\t\tnorm = mpl.colors.Normalize(vmin=0, vmax=3000)\n",
    "\t\trng = np.arange(267.5,270.5,0.25)\n",
    "\t\tt2_contourf = m.contourf(x,y,to_np(t2[:,:]),cmap=cmocean.cm.thermal,levels=rng,extend='both',zorder=2)\n",
    "\t\tw_contours = m.contour(x,y,to_np(wi[:,:]),colors='magenta',levels=[0.25],zorder=3)\n",
    "\t\tm.contour(x,y,to_np(seaice),levels=[0.0001,0.25,0.5,0.75,1.0],colors=['magenta','pink','pink','pink'],linewidths=2.,zorder=4)\n",
    "\n",
    "\t\t# Add geographic outlines\n",
    "\t\tm.drawcoastlines(linewidth=2.5,zorder=3)\n",
    "\t\tif j == 0:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[True,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,True],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\telse:\n",
    "\t\t\tm.drawparallels(np.arange(60.,94.,0.5),labels=[False,False,False,False],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\t\tm.drawmeridians(np.arange(-40.,70.,1.5),labels=[False,False,False,True],fontsize=22,linewidth=3.0,zorder=3)\n",
    "\t\tm.drawcountries(linewidth=2.0,zorder=3)\n",
    "\t\tm.fillcontinents(color='coral',lake_color='coral',zorder=3)\n",
    "        \n",
    "\t\tx_anx,y_anx = m(15.874211,69.232286)\n",
    "\t\tm.scatter(x_anx,y_anx,marker='D',s=360,facecolor='yellow',edgecolor='k',zorder=8)\n",
    "        \n",
    "\t\tif j == 0 or j == 2:\n",
    "\t\t\tif j == 0:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.5)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.5)\n",
    "\t\t\telse:\n",
    "\t\t\t\tx1,y1 = m(12.5,69.4)\n",
    "\t\t\t\tx2,y2 = m(16.5,69.4)\n",
    "\t\t\tm.plot([x1,x2],[y1,y2],c='limegreen',lw=3.,ls='--',zorder=8)\n",
    "                \n",
    "\t\tif j == 3:\n",
    "\t\t\tplt.tight_layout()\n",
    "\t\t\tplt.subplots_adjust(right=0.90)\n",
    "\t\t\tcbar = fig.colorbar(w_contourf, cax=ax1)\n",
    "\t\t\tcbar.ax.tick_params(labelsize=18)\n",
    "\t\t\tcbar.ax.set_ylabel('$W$ [m s$^{-1}$]', labelpad=40, rotation=270, size=24)\n",
    "\t\t\tcbar = fig.colorbar(conv_contourf, cax=ax2)\n",
    "\t\t\tcbar.ax.tick_params(labelsize=18)\n",
    "\t\t\tcbar.ax.set_ylabel('Divergence [s$^{-1}$]', labelpad=40, rotation=270, size=24)\n",
    "\t\t\tcbar = fig.colorbar(t2_contourf, cax=ax3)\n",
    "\t\t\tcbar.ax.tick_params(labelsize=18)\n",
    "\t\t\tcbar.ax.set_ylabel('$T_{2}$ [K]', labelpad=40, rotation=270, size=24)\n",
    "            \n",
    "\n",
    "    \n",
    "\tplt.savefig(savedir + \"compare_mar\" + day + \"_mynn_edmf_zoom_more_\" + region + \"_\" + active_dom + \"_PANEL_0p5zi_031320_1800.png\",dpi=200,bbox_inches='tight')\n",
    "\tplt.close()\n",
    "    \n",
    "\ts1.close()\n",
    "\ts3.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
